{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "279977fbd35ba1d4",
   "metadata": {},
   "source": [
    "# Streaming With LangChain\n",
    "Streaming is critical in making applications based on LLMs feel responsive to end-users.<br>\n",
    "流式处理对于使基于感觉LLMs的应用程序对最终用户的响应至关重要。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc56406b817c5a65",
   "metadata": {},
   "source": [
    "## Using Stream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f99e7737bc05ebbe",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-23T04:20:34.868686Z",
     "start_time": "2024-07-23T04:20:31.278899Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|Hello|!| I| am| a| language| model| AI| created| by| Open|AI|.| I| am| constantly| learning| and| improving| my| abilities| to| assist| users| with| a| wide| range| of| tasks| and| information|.| I| am| here| to| help| answer| any| questions| you| may| have| or| engage| in| conversation| on| various| topics|.| How| can| I| assist| you| today|?|||"
     ]
    }
   ],
   "source": [
    "# Showing the example using anthropic, but you can use\n",
    "# your favorite chat model!\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "model = ChatOpenAI()\n",
    "\n",
    "chunks = []\n",
    "async for chunk in model.astream(\"hello. tell me something about yourself\"):\n",
    "    chunks.append(chunk)\n",
    "    print(chunk.content, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "df9d600b7ddd8d0f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-23T04:23:38.656220Z",
     "start_time": "2024-07-23T04:23:14.583744Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|Why| did| the| par|rot| go| to| the| doctor|?| Because| it| was| feeling| a| little| \"|ill|-e|agle|\"|!|||"
     ]
    }
   ],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(\"tell me a joke about {topic}\")\n",
    "parser = StrOutputParser()\n",
    "chain = prompt | model | parser\n",
    "\n",
    "async for chunk in chain.astream({\"topic\": \"parrot\"}):\n",
    "    print(chunk, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7b8e1622",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{}\n",
      "{'countries': []}\n",
      "{'countries': [{}]}\n",
      "{'countries': [{'name': ''}]}\n",
      "{'countries': [{'name': 'France'}]}\n",
      "{'countries': [{'name': 'France', 'population': 670}]}\n",
      "{'countries': [{'name': 'France', 'population': 670760}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': ''}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain'}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 473}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 473299}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 47329981}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 47329981}, {}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 47329981}, {'name': ''}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 47329981}, {'name': 'Japan'}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 47329981}, {'name': 'Japan', 'population': 126}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 47329981}, {'name': 'Japan', 'population': 126476}]}\n",
      "{'countries': [{'name': 'France', 'population': 67076000}, {'name': 'Spain', 'population': 47329981}, {'name': 'Japan', 'population': 126476461}]}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.output_parsers import JsonOutputParser\n",
    "\n",
    "chain = (\n",
    "    model | JsonOutputParser()\n",
    ")  # Due to a bug in older versions of Langchain, JsonOutputParser did not stream results from some models\n",
    "async for text in chain.astream(\n",
    "    'output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of \"countries\" which contains a list of countries. Each country should have the key `name` and `population`'\n",
    "):\n",
    "    print(text, flush=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "05a3749c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['France', 'Spain', 'Japan']|"
     ]
    }
   ],
   "source": [
    "from langchain_core.output_parsers import (\n",
    "    JsonOutputParser,\n",
    ")\n",
    "\n",
    "\n",
    "# A function that operates on finalized inputs\n",
    "# rather than on an input_stream\n",
    "def _extract_country_names(inputs):\n",
    "    \"\"\"A function that does not operates on input streams and breaks streaming.\"\"\"\n",
    "    if not isinstance(inputs, dict):\n",
    "        return \"\"\n",
    "\n",
    "    if \"countries\" not in inputs:\n",
    "        return \"\"\n",
    "\n",
    "    countries = inputs[\"countries\"]\n",
    "\n",
    "    if not isinstance(countries, list):\n",
    "        return \"\"\n",
    "\n",
    "    country_names = [\n",
    "        country.get(\"name\") for country in countries if isinstance(country, dict)\n",
    "    ]\n",
    "    return country_names\n",
    "\n",
    "\n",
    "chain = model | JsonOutputParser() | _extract_country_names\n",
    "\n",
    "async for text in chain.astream(\n",
    "    'output a list of the countries france, spain and japan and their populations in JSON format. Use a dict with an outer key of \"countries\" which contains a list of countries. Each country should have the key `name` and `population`'\n",
    "):\n",
    "    print(text, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08072533",
   "metadata": {},
   "source": [
    "### Non-streaming components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "610730d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[Document(page_content='harrison worked at kensho'),\n",
       "  Document(page_content='harrison likes spicy food')]]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "template = \"\"\"Answer the question based only on the following context:\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "vectorstore = FAISS.from_texts(\n",
    "    [\"harrison worked at kensho\", \"harrison likes spicy food\"],\n",
    "    embedding=OpenAIEmbeddings(),\n",
    ")\n",
    "retriever = vectorstore.as_retriever()\n",
    "\n",
    "chunks = [chunk for chunk in retriever.stream(\"where did harrison work?\")]\n",
    "chunks"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
